Skip to content

[PyTorch] Enable NVFP4 row-scaled (per-token) backward for GroupedLinear - #3324

Open
cael-ling wants to merge 3 commits into
NVIDIA:mainfrom
cael-ling:pr/nvfp4-row-scaled-moe-backward
Open

[PyTorch] Enable NVFP4 row-scaled (per-token) backward for GroupedLinear#3324
cael-ling wants to merge 3 commits into
NVIDIA:mainfrom
cael-ling:pr/nvfp4-row-scaled-moe-backward

Conversation

@cael-ling

@cael-ling cael-ling commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Description

Extend the row-scaled NVFP4 support added for dense Linear (originally #2931 and #3206) to the MoE GroupedLinear module, so the wgrad is computed in NVFP4 instead of falling back to high precision.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@cael-ling
[PyTorch] Enable NVFP4 row-scaled (per-token) backward for GroupedLinear
e1bb076 
Extend the row-scaled NVFP4 support added for dense Linear to the
MoE GroupedLinear module, so the wgrad is computed in NVFP4 instead of
falling back to high precision.

Signed-off-by: Cael Ling <caell@nvidia.com>
@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Aug 7, 2026
@greptile-apps

greptile-apps Bot commented Aug 7, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR enables row-scaled NVFP4 GroupedLinear backward by routing affected grouped GEMMs through dense per-expert operations and adding the required columnwise amax allocation.

  • Supports row-scaled NVFP4 tensors in either grouped-GEMM operand.
  • Allocates and exposes per-column amax metadata for transposed row-scaled tensors.
  • Adds bit-exact GroupedLinear-versus-Linear backward coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
transformer_engine/pytorch/cpp_extensions/gemm.py Routes grouped GEMMs containing row-scaled NVFP4 operands through compatible per-expert dense GEMMs.
transformer_engine/pytorch/csrc/extensions/cast.cpp Enables columnwise storage for row-scaled NVFP4 bulk allocation and assigns per-column amax shapes.
transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh Makes the general columnwise amax reduction available without an unnecessary SM100 device guard.
tests/pytorch/test_grouped_linear.py Adds bit-exact forward, dgrad, and wgrad comparisons against independent dense Linear layers.

Sequence Diagram

sequenceDiagram
  participant GL as GroupedLinear
  participant Q as NVFP4 quantization
  participant GG as general_grouped_gemm
  participant DG as general_gemm
  participant K as CUDA kernels
  GL->>Q: Quantize per-expert activations and gradients
  Q->>Q: Allocate rowwise and columnwise amax metadata
  GL->>GG: Submit grouped forward/dgrad/wgrad
  loop Each expert
    GG->>DG: Invoke dense row-scaled GEMM
    DG->>K: Quantize, transpose, and multiply
    K-->>DG: Expert output
  end
  GG-->>GL: Combined outputs and gradients
Loading

Reviews (2): Last reviewed commit: "Remove unhelpful comments" | Re-trigger Greptile

timmoon10
timmoon10 reviewed Aug 11, 2026
View reviewed changes
Comment thread transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh Outdated
Comment thread transformer_engine/pytorch/csrc/extensions/cast.cpp Outdated
@timmoon10
Remove unhelpful comments
88097b2 
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
timmoon10
timmoon10 approved these changes Aug 11, 2026
View reviewed changes

@timmoon10 timmoon10 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, pending CI.

@timmoon10

timmoon10 commented Aug 11, 2026

Copy link
Copy Markdown
Member

/te-ci

@ptrendx ptrendx self-assigned this Aug 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10 timmoon10 timmoon10 approved these changes
@Oleg-Goncharov Oleg-Goncharov Awaiting requested review from Oleg-Goncharov Oleg-Goncharov is a code owner
@ksivaman ksivaman Awaiting requested review from ksivaman ksivaman is a code owner

Assignees

@ptrendx ptrendx

Labels

community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@cael-ling @timmoon10 @ptrendx